/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2024-2024. All rights reserved.
 * Description: GPU/NPU send/recv typical interfaces
 */

/*
 * 该特性代码不涉及开源
 */

#include "interface_hccl.h"
#include "typical_param_check.h"
#include "dltdt_function.h"
#include "adapter_tdt.h"
#include "hccl_common.h"
#include "network_manager_pub.h"
#include "externalinput.h"
#include "rdma_resource_manager.h"
#include "typical_mr_manager.h"
#include "typical_sync_mem.h"
#include "typical_window_mem.h"
#include "typical_qp_manager.h"
#include "send_recv_executor.h"

using namespace hccl;
constexpr u32 DEVISOR_VALUE_FOUR = 4;

struct mr_info AscendMrInfo2MrInfo(AscendMrInfo* ascendMrInfo)
{
    struct mr_info innerMrInfo;
    innerMrInfo.addr = reinterpret_cast<void*>(ascendMrInfo->addr);
    innerMrInfo.size = ascendMrInfo->size;
    innerMrInfo.lkey = ascendMrInfo->key;
    return innerMrInfo;
}

HcclResult hcclCreateAscendQP(AscendQPInfo* ascendQPInfo)
{
    CHK_PTR_NULL(ascendQPInfo);
    struct typical_qp qpInfo;
    CHK_RET(TypicalQpManager::GetInstance().CreateQp(qpInfo));
    ascendQPInfo->qpn = qpInfo.qpn;
    ascendQPInfo->gid_idx = qpInfo.gid_idx;
    for (uint32_t i = 0; i < GID_LENGTH; i++) {
        ascendQPInfo->gid[i] = qpInfo.gid[i];
    }
    ascendQPInfo->psn = qpInfo.psn;
    HCCL_INFO("hcclCreateAscendQP success! qpn %u, gid index %u, psn %u", ascendQPInfo->qpn,
        ascendQPInfo->gid_idx, ascendQPInfo->psn);

    return HCCL_SUCCESS;
}

HcclResult hcclModifyAscendQP(AscendQPInfo* localQPInfo, AscendQPInfo* remoteQPInfo)
{
    AscendQPQos qpQos;
    qpQos.sl = GetExternalInputRdmaServerLevel();
    qpQos.tc = GetExternalInputRdmaTrafficClass();
    CHK_RET(hcclModifyAscendQPEx(localQPInfo, remoteQPInfo, &qpQos));
    return HCCL_SUCCESS;
}

HcclResult hcclModifyAscendQPEx(AscendQPInfo* localQPInfo, AscendQPInfo* remoteQPInfo, AscendQPQos* qpQos)
{
    CHK_PTR_NULL(localQPInfo);
    CHK_PTR_NULL(remoteQPInfo);
    CHK_PTR_NULL(qpQos);
    CHK_PRT_RET((qpQos->sl < HCCL_RDMA_SL_MIN || qpQos->sl > HCCL_RDMA_SL_MAX),
        HCCL_ERROR("[hcclModifyAscendQPEx]The value of sl[%u] is invalid. except: [%u, %u].",
        qpQos->sl, HCCL_RDMA_SL_MIN, HCCL_RDMA_SL_MAX), HCCL_E_PARA);
    CHK_PRT_RET((qpQos->tc < HCCL_RDMA_TC_MIN || qpQos->tc > HCCL_RDMA_TC_MAX),
        HCCL_ERROR("[hcclModifyAscendQPEx]The value of tc[%u] is invalid. except: [%u, %u].",
        qpQos->tc, HCCL_RDMA_TC_MIN, HCCL_RDMA_TC_MAX), HCCL_E_PARA);

    // 设置的RDMATrafficClass需要是4的整数倍, 否则报错
    CHK_PRT_RET(qpQos->tc % DEVISOR_VALUE_FOUR != 0,
        HCCL_ERROR("[hcclModifyAscendQPEx]The value of tc[%u] is not a multiple of 4.",
        qpQos->tc), HCCL_E_PARA);    

    struct typical_qp localQp;
    localQp.qpn = localQPInfo->qpn;
    localQp.gid_idx = localQPInfo->gid_idx;
    for (uint32_t i = 0; i < GID_LENGTH; i++) {
        localQp.gid[i] = localQPInfo->gid[i];
    }
    localQp.psn = localQPInfo->psn;
    localQp.sl = qpQos->sl;
    localQp.tc = qpQos->tc;

    struct typical_qp remoteQp;
    remoteQp.qpn = remoteQPInfo->qpn;
    remoteQp.gid_idx = remoteQPInfo->gid_idx;
    for (uint32_t i = 0; i < GID_LENGTH; i++) {
        remoteQp.gid[i] = remoteQPInfo->gid[i];
    }
    remoteQp.psn = remoteQPInfo->psn;
    CHK_RET(TypicalQpManager::GetInstance().ModifyQp(localQp, remoteQp));
    return HCCL_SUCCESS;
}

HcclResult hcclDestroyAscendQP(AscendQPInfo* ascendQPInfo)
{
    CHK_PTR_NULL(ascendQPInfo);
    struct typical_qp qpInfo;
    qpInfo.qpn = ascendQPInfo->qpn;
    qpInfo.gid_idx = ascendQPInfo->gid_idx;
    for (uint32_t i = 0; i < GID_LENGTH; i++) {
        qpInfo.gid[i] = ascendQPInfo->gid[i];
    }
    qpInfo.psn = ascendQPInfo->psn;
    CHK_RET(TypicalQpManager::GetInstance().DestroyQp(qpInfo));
    return HCCL_SUCCESS;
}

HcclResult hcclAllocWindowMem(void **ptr, size_t len)
{
    return TypicalWindowMem::GetInstance().AllocWindowMem(ptr, len);
}

HcclResult hcclFreeWindowMem(void *ptr)
{
    return TypicalWindowMem::GetInstance().FreeWindowMem(ptr);
}

HcclResult hcclAllocSyncMem(int32_t **ptr)
{
    return TypicalSyncMem::GetInstance().AllocSyncMem(ptr);
}

HcclResult hcclFreeSyncMem(int32_t *ptr)
{
    return TypicalSyncMem::GetInstance().FreeSyncMem(ptr);
}

HcclResult hcclRegisterMem(AscendMrInfo* memInfo)
{
    CHK_PTR_NULL(memInfo);
    mr_info mrInfo = {};
    mrInfo.addr = reinterpret_cast<void *>(static_cast<uintptr_t>(memInfo->addr));
    mrInfo.size = memInfo->size;
    mrInfo.access = RA_ACCESS_REMOTE_WRITE | RA_ACCESS_LOCAL_WRITE | RA_ACCESS_REMOTE_READ;
    CHK_RET(TypicalMrManager::GetInstance().RegisterMem(mrInfo));
    memInfo->key = mrInfo.lkey;
    return HCCL_SUCCESS;
}

HcclResult hcclDeRegisterMem(AscendMrInfo* memInfo)
{
    CHK_PTR_NULL(memInfo);
    mr_info mrInfo = {};
    mrInfo.addr = reinterpret_cast<void *>(static_cast<uintptr_t>(memInfo->addr));
    mrInfo.size = memInfo->size;
    mrInfo.access = RA_ACCESS_REMOTE_WRITE | RA_ACCESS_LOCAL_WRITE | RA_ACCESS_REMOTE_READ;
    mrInfo.lkey = memInfo->key;
    CHK_RET(TypicalMrManager::GetInstance().DeRegisterMem(mrInfo));
    return HCCL_SUCCESS;
}

HcclResult HcclSendByAscendQP(void* sendBuf, uint64_t count, HcclDataType dataType,
    AscendSendRecvInfo* sendRecvInfo, aclrtStream stream)
{
    CHK_RET(CheckParam(sendBuf, count, dataType, sendRecvInfo, stream));
    QpHandle qpHandle;
    CHK_RET(TypicalQpManager::GetInstance().GetQpHandleByQpn(sendRecvInfo->localQPinfo->qpn, qpHandle));
    SendRecvExecutor executor(stream, qpHandle, AscendMrInfo2MrInfo(sendRecvInfo->localWindowMem),
                                                    AscendMrInfo2MrInfo(sendRecvInfo->remoteWindowMem),
                                                    AscendMrInfo2MrInfo(sendRecvInfo->localSyncMemPrepare),
                                                    AscendMrInfo2MrInfo(sendRecvInfo->localSyncMemDone),
                                                    AscendMrInfo2MrInfo(sendRecvInfo->localSyncMemAck),
                                                    AscendMrInfo2MrInfo(sendRecvInfo->remoteSyncMemPrepare),
                                                    AscendMrInfo2MrInfo(sendRecvInfo->remoteSyncMemDone),
                                                    AscendMrInfo2MrInfo(sendRecvInfo->remoteSyncMemAck),
                                                    sendRecvInfo->immData);
    CHK_RET(executor.Init());
    CHK_RET(executor.Send(sendBuf, count, dataType));
    return HCCL_SUCCESS;
}

HcclResult HcclRecvByAscendQP(void* recvBuf, uint64_t count, HcclDataType dataType,
    AscendSendRecvInfo* sendRecvInfo, aclrtStream stream)
{
    CHK_RET(CheckParam(recvBuf, count, dataType, sendRecvInfo, stream));
    QpHandle qpHandle;
    CHK_RET(TypicalQpManager::GetInstance().GetQpHandleByQpn(sendRecvInfo->localQPinfo->qpn, qpHandle));
    SendRecvExecutor executor(stream, qpHandle, AscendMrInfo2MrInfo(sendRecvInfo->localWindowMem),
                                                    AscendMrInfo2MrInfo(sendRecvInfo->remoteWindowMem),
                                                    AscendMrInfo2MrInfo(sendRecvInfo->localSyncMemPrepare),
                                                    AscendMrInfo2MrInfo(sendRecvInfo->localSyncMemDone),
                                                    AscendMrInfo2MrInfo(sendRecvInfo->localSyncMemAck),
                                                    AscendMrInfo2MrInfo(sendRecvInfo->remoteSyncMemPrepare),
                                                    AscendMrInfo2MrInfo(sendRecvInfo->remoteSyncMemDone),
                                                    AscendMrInfo2MrInfo(sendRecvInfo->remoteSyncMemAck),
                                                    sendRecvInfo->immData);
    CHK_RET(executor.Init());
    CHK_RET(executor.Receive(recvBuf, count, dataType));
    return HCCL_SUCCESS;
}


HcclResult hcclAscendRdmaInit()
{
    CHK_RET(RdmaResourceManager::GetInstance().Init());
    return HCCL_SUCCESS;
}
HcclResult hcclAscendRdmaDeInit()
{
    CHK_RET(RdmaResourceManager::GetInstance().DeInit());
    return HCCL_SUCCESS;
}

HcclResult HcclGetCqeErrInfoList(struct HcclErrCqeInfo *infoList, uint32_t *num)
{
    CHK_PTR_NULL(infoList);
    CHK_PTR_NULL(num);
    u32 arrLen = *num;
    struct cqe_err_info errCqeList[arrLen];
    CHK_RET(RdmaResourceManager::GetInstance().GetCqeErrInfo(errCqeList, num));

    CHK_PRT_RET(*num > arrLen, HCCL_ERROR("[HcclGetCqeErrInfoList] GetCqeErrInfo num[%u] is larger than "
        "infoList user given[%u].", num, arrLen), HCCL_E_INTERNAL);
    for (u32 i = 0; i < *num; i++) {
        infoList[i].status = errCqeList[i].status;
        infoList[i].qpn = errCqeList[i].qpn;
        infoList[i].time = errCqeList[i].time;
        time_t tmpt = static_cast<time_t>(errCqeList[i].time.tv_sec);
        struct tm errTime;
        localtime_r(&tmpt, &errTime);
        HCCL_INFO("[HcclGetCqeErrInfoList] Error Cqe status[%d], qpn[%d], time[%04u-%02d-%02d %02d:%0d:%02d.%06u]", 
            errCqeList[i].status, errCqeList[i].qpn, errTime.tm_year + TIME_FROM_1900,
            errTime.tm_mon + 1,
            errTime.tm_mday,
            errTime.tm_hour,
            errTime.tm_min,
            errTime.tm_sec,
            static_cast<u32>(errCqeList[i].time.tv_usec));
    }
    return HCCL_SUCCESS;
}

HcclResult HcclGetCqeErrInfoListByQpn(uint32_t qpn, struct HcclErrCqeInfo *infoList, uint32_t *num)
{
    CHK_PTR_NULL(infoList);
    CHK_PTR_NULL(num);
    CHK_RET(RdmaResourceManager::GetInstance().GetCqeErrInfoByQpn(qpn, infoList, num));
    return HCCL_SUCCESS;
}